#include "dag_graph_operator.h"
#include "dag_apply_link.h"
#include "dag_types.h"
#include <string>

using namespace std;
using namespace dag;

namespace {
	
	/* delete all object pointers within a set and clear set */
	template <typename _obj>
	void clear_set(set<_obj*>& objSet) {
		typename set<_obj*>::iterator objIter = objSet.begin();
		while (objIter != objSet.end()) {
			delete *objIter;
			objIter++;
		}
		objSet.clear();
	}
	
	/* grab the next object by looking at the id in the binary stream */
	template <typename _obj>
	_obj* get_next_object(istream& binaryStream,
						  const map<dag_int32_t, _obj*>& idMap) {
		dag_int32_t id;
		binaryStream.read((char*)&id, dag_int32_size);
		return get_by_id(idMap, id);
	}
	
	/* get the next object by id */
	template <typename _obj>
	_obj* get_by_id(const map<dag_int32_t, _obj*>& idMap, dag_int32_t id) {
		if (idMap.count(id) == 0) {
			throw SerialError("corrupt stream");
		}
		return idMap.find(id)->second;
	}
	
	/* serialize an entire set of objects.  First store their id, then the
	 * object itself. */
	template <typename _obj>
	void serialize_set(ostream& binaryStream,
					   const map<const _obj*, dag_int32_t>& idmap,
					   const set<_obj*>& objSet) {
		dag_uint32_t numObjs = objSet.size();
		binaryStream.write((const char*)&numObjs, dag_uint32_size);

		typename set<_obj*>::const_iterator objIter = objSet.begin();
		while (objIter != objSet.end()) {
			if (idmap.count(*objIter) == 0) {
				throw DagError("missing object id");
			}
			
			/* write internal id */
			binaryStream.write((const char*)&idmap.find(*objIter)->second,
							   dag_int32_size);
			/* write object */
			(*objIter)->serialize(binaryStream);
			objIter++;
		}
	}
	
	/* load the operator / frame mapping */
	void load_op_frame_map(istream& binaryStream,
						   const map<dag_int32_t, FrameOperator*>& frameOpIds,
						   const map<dag_int32_t, Frame*>& frameIds,
						   map<FrameOperator*, Frame*>& opFrames) {
		
		dag_uint32_t count;
		dag_int32_t id;
		binaryStream.read((char*)&count, dag_uint32_size);
		for (dag_uint32_t i = 0; i < count; i++) {
			binaryStream.read((char*)&id, dag_int32_size);
			FrameOperator* op = get_by_id(frameOpIds, id);
			binaryStream.read((char*)&id, dag_int32_size);
			Frame* input = get_by_id(frameIds, id);
			opFrames[op] = input;
		}
	}

	void _visit_node_(const Node* node,
					  set<const Node*>& visited,
					  const Graph& graph, 
					  list<const Node*>& sortedNodes) {
		if (visited.count(node) == 0) {
			visited.insert(node);
			vector<Edge*> inputs = graph.getInputEdges(node);
			for (size_t i = 0; i < inputs.size(); i++) {
				if (inputs[i] == NULL) {
					throw NullPtr();
				}
				_visit_node_(inputs[i]->from,
							 visited,
							 graph,
							 sortedNodes);
			}
			sortedNodes.push_back(node);
		}
	}
			
	void topological_sort(const Graph& graph,
						  list<const Node*>& sortedNodes) {
		set<const Node*> visited;
		set<Node*>::const_iterator nodeIter = graph.getNodes().begin();
		while (nodeIter != graph.getNodes().end()) {
			_visit_node_(*nodeIter, 
						 visited, 
						 graph, 
						 sortedNodes);
			nodeIter++;
		}
	}
	
}


string GraphOperator::m_OpId("graph_operator");

GraphOperator::GraphOperator() {
}

GraphOperator::GraphOperator(const GraphNode& graphNode, bool reverse) {
	if (reverse) {
		initReverse(graphNode);
	}
	else {
		initForward(graphNode);
	}
}

//GraphOperator::GraphOperator(const GraphOperator& copy) {
//	/* TODO */
//}

GraphOperator::~GraphOperator() throw() {
	try {
		clear();
	}
	catch (...) {
	}
}

void GraphOperator::operate(const Frame& input, Frame& output) {
	
	FrameOperator* frameOp = NULL;
	const Frame* inputFrame = NULL;
	Frame* outputFrame = NULL;
	
	/* apply frame operators */
	for (size_t i = 0; i < m_OrderedFrameOps.size(); i++) {
		
		frameOp = m_OrderedFrameOps[i];
		
		if (i == 0) {
			/* head operator */
			inputFrame = inputFrame;
			outputFrame = m_OpOutputs[frameOp];	
		}
		else if (i == (m_OrderedFrameOps.size() - 1)) {
			/* tail operator */
			inputFrame = m_OpInputs[frameOp];
			outputFrame = &output;
		}
		else {
			/* middle operators */
			inputFrame = m_OpInputs[frameOp];
			outputFrame = m_OpOutputs[frameOp];
		}
		
		frameOp->operate(*inputFrame, *outputFrame);
		
		/* apply output of frame operator to other input frames */
		const vector<FrameAndLink>* links = m_OpOutputLinks[frameOp];
		if (links != NULL) {
			for (size_t s = 0; s < links->size(); s++) {
				/* TODO: optimize this by using link's asArray() 
				 * functionality */
				apply_link<true>(*outputFrame,
								 *links->at(s).link,
								 *links->at(s).frame);
			}
		}		
	}
}

auto_ptr<FrameOperator> GraphOperator::inflate(istream& binaryStream) {
	
	auto_ptr<GraphOperator> graphOp(new GraphOperator());
	map<dag_int32_t, FrameOperator*> frameOpIds;
	map<dag_int32_t, Frame*> frameIds;
	map<dag_int32_t, FrameLink*> linkIds;
	
	/* TODO: consider more specific serial error.  perhaps an "FailedInflation"
	 * or "CorruptStream" error */
	
	/* load operators, frames and links */
	dag_int32_t id = 0;
	dag_uint32_t count = 0;
	
	/* load operators */
	binaryStream.read((char*)&count, dag_uint32_size);
	for (dag_uint32_t i = 0; i < count; i++) {
		binaryStream.read((char*)&id, dag_int32_size);
		auto_ptr<FrameOperator> newFrameOp;
		newFrameOp = FrameOperator::inflate(binaryStream);
		if (newFrameOp.get() == NULL) {
			throw SerialError("failed to inflate frame operator");
		}
		
		frameOpIds[id] = newFrameOp.get();
		graphOp->m_FrameOps.insert(newFrameOp.release());
	}
	
	/* load frames */
	binaryStream.read((char*)&count, dag_uint32_size);
	for (dag_uint32_t i = 0; i < count; i++) {
		binaryStream.read((char*)&id, dag_int32_size);
		auto_ptr<FrameSpec> newFrameSpec = FrameSpec::inflate(binaryStream);
		if (newFrameSpec.get() == NULL) {
			throw SerialError("failed to inflate frame spec");
		}
		
		auto_ptr<Frame> newFrame(new Frame(*newFrameSpec));
		
		frameIds[id] = newFrame.get();
		graphOp->m_Frames.insert(newFrame.release());
	}
	
	
	/* load links */
	binaryStream.read((char*)&count, dag_uint32_size);
	for (dag_uint32_t i = 0; i < count; i++) {
		binaryStream.read((char*)&id, dag_int32_size);
		auto_ptr<FrameLink> newLink(NULL);
		newLink.reset(new FrameLink(FrameLink::inflate(binaryStream)));
		if (newLink.get() == NULL) {
			throw SerialError("failed to inflate frame link");
		}
		
		linkIds[id] = newLink.get();
		graphOp->m_Links.insert(newLink.release());
	}
	
	/* load operator order */
	binaryStream.read((char*)&count, dag_uint32_size);
	for (dag_uint32_t i = 0; i < count; i++) {
		graphOp->m_OrderedFrameOps.push_back(get_next_object(binaryStream,
															 frameOpIds));
	}
	
	/* lodad operator inputs */
	load_op_frame_map(binaryStream,
					  frameOpIds,
					  frameIds,
					  graphOp->m_OpInputs);
	
	/* load operator outputs */
	load_op_frame_map(binaryStream,
					  frameOpIds,
					  frameIds, 
					  graphOp->m_OpOutputs);
	
	/* load link mapping */
	binaryStream.read((char*)&count, dag_uint32_size);
	for (dag_uint32_t i = 0; i < count; i++) {
		FrameOperator* frameOp = get_next_object(binaryStream, frameOpIds);

		dag_uint32_t linkCount = 0;
		binaryStream.read((char*)&linkCount, dag_uint32_size);
		vector<FrameAndLink> fands;
		
		for (dag_uint32_t j = 0; j < linkCount; j++) {
			FrameAndLink fs;
			fs.link = get_next_object(binaryStream, linkIds);
			fs.frame = get_next_object(binaryStream, frameIds);
			fands.push_back(fs);
		}
		
		vector<FrameAndLink>* pv = new vector<FrameAndLink>(fands);
		graphOp->m_LinkVectors.insert(pv);
		graphOp->m_OpOutputLinks[frameOp] = pv;
	}
	
	return auto_ptr<FrameOperator>(graphOp);
}

void GraphOperator::clear() {
	/* only delete objects in sets.  all other containers may have ambiguities.
	 * the sets exist to keep unique copies of objects */
	
	/* clear_set(...) deletes all objects within the set and empties the set */
	clear_set(m_FrameOps);
	clear_set(m_Frames);
	clear_set(m_Links);
	clear_set(m_LinkVectors);
	
	m_OpInputs.clear();
	m_OpOutputs.clear();
	m_OpOutputLinks.clear(); 
	m_OrderedFrameOps.clear();
}

const string& GraphOperator::_frameOperatorID() const {
	return m_OpId;
}

void GraphOperator::_serialize(ostream& binaryStream) const {
	/* assign each object an id, so we can use id's to rebuild links */
	dag_int32_t itemId = 1;
	map<const FrameOperator*, dag_int32_t> frameOpIds;
	map<const Frame*, dag_int32_t> frameIds;
	map<const FrameLink*, dag_int32_t> linkIds;
	
	/* build frame op ids */
	set<FrameOperator*>::const_iterator opIter = m_FrameOps.begin();
	while (opIter != m_FrameOps.end()) {
		frameOpIds[*opIter] = itemId++;
		opIter++;
	}
	
	/* build frame ids */
	set<Frame*>::const_iterator frameIter = m_Frames.begin();
	while (frameIter != m_Frames.end()) {
		frameIds[*frameIter] = itemId++;
		frameIter++;
	}
	
	set<FrameLink*>::const_iterator linkIter = m_Links.begin();
	while (linkIter != m_Links.end()) {
		linkIds[*linkIter] = itemId++;
		linkIter++;
	}
	
	/* write ops */
	serialize_set(binaryStream, frameOpIds, m_FrameOps);
	
	/* write frames */
	serialize_set(binaryStream, frameIds, m_Frames);
	
	/* write links */
	serialize_set(binaryStream, linkIds, m_Links);
	
	/* write op order */
	dag_uint32_t numOps = m_OrderedFrameOps.size();
	binaryStream.write((const char*)&numOps, dag_uint32_size);
	
	for (size_t i = 0; i < m_OrderedFrameOps.size(); i++) {
		itemId = frameOpIds[m_OrderedFrameOps[i]];
		binaryStream.write((const char*)&itemId, dag_int32_size);
	}
	
	
	/* write op inputs */
	numOps = m_OpInputs.size();
	binaryStream.write((const char*)&numOps, dag_uint32_size);
	
	map<FrameOperator*, Frame*>::const_iterator inputIter = m_OpInputs.begin();
	while (inputIter != m_OpInputs.end()) {
		itemId = frameOpIds[inputIter->first];
		binaryStream.write((const char*)&itemId, dag_int32_size);
		itemId = frameIds[inputIter->second];
		binaryStream.write((const char*)&itemId, dag_int32_size);
		inputIter++;
	}
	
	/* write op outputs */
	numOps = m_OpOutputs.size();
	binaryStream.write((const char*)&numOps, dag_uint32_size);
	
	map<FrameOperator*, Frame*>::const_iterator outIter = m_OpOutputs.begin();
	while (outIter != m_OpOutputs.end()) {
		itemId = frameOpIds[outIter->first];
		binaryStream.write((const char*)&itemId, dag_int32_size);
		itemId = frameIds[outIter->second];
		binaryStream.write((const char*)&itemId, dag_int32_size);
		outIter++;
	}
	
	/* write out links */
	numOps = m_OpOutputLinks.size();
	binaryStream.write((const char*)&numOps, dag_uint32_size);
	
	map<FrameOperator*, std::vector<FrameAndLink>*>::const_iterator linkVecIter;
	linkVecIter = m_OpOutputLinks.begin();
	while (linkVecIter != m_OpOutputLinks.end()) {
		itemId = frameOpIds[linkVecIter->first];
		binaryStream.write((const char*)&itemId, dag_int32_size);
		const vector<FrameAndLink>& fas = *(linkVecIter->second);
		dag_uint32_t size = fas.size();
		binaryStream.write((const char*)&size, dag_uint32_size);
		
		for (size_t i = 0; i < fas.size(); i++) {
			itemId = linkIds[fas[i].link];
			binaryStream.write((const char*)&itemId, dag_int32_size);
			itemId = frameIds[fas[i].frame];
			binaryStream.write((const char*)&itemId, dag_int32_size);
		}
	}
}

void GraphOperator::initForward(const GraphNode& graphNode) {
	if (graphNode.getGraph() == NULL) {
		/* Todo: better error reporting */
		throw NullPtr();
	}
	const Graph& graph = *graphNode.getGraph();
	
	/* sort nodes in graph */
	list<const Node*> sortedNodes;
	topological_sort(graph, sortedNodes);
	
	/* add operators and i/o frames from graph in order */
	map<const Node*, FrameOperator*> nodeOpMap;
	list<const Node*>::iterator nodeIter = sortedNodes.begin();
	
	while (nodeIter != sortedNodes.end()) {
		auto_ptr<FrameOperator> nodeFrameOp(NULL);
		auto_ptr<Frame> inputFrame(NULL);
		auto_ptr<Frame> outputFrame(NULL);
		
		nodeFrameOp = (*nodeIter)->createForwardOp();
		inputFrame.reset(new Frame((*nodeIter)->inputSpec()));
		outputFrame.reset(new Frame((*nodeIter)->outputSpec()));
		if (nodeFrameOp.get() == NULL) {
			/* TODO: better error reporting */
			throw NullPtr();
		}
		
		nodeOpMap[*nodeIter] = nodeFrameOp.get();
		
		m_OrderedFrameOps.push_back(nodeFrameOp.get());
		m_OpInputs[nodeFrameOp.get()] = inputFrame.get();
		m_OpOutputs[nodeFrameOp.get()] = outputFrame.get();
		
		m_FrameOps.insert(nodeFrameOp.release());
		m_Frames.insert(inputFrame.release());
		m_Frames.insert(outputFrame.release());
		nodeIter++;
	}
	
	/* take a second pass and add links. we couldn't add the links originally
	 * because they require all the nodes to be loaded. */
	nodeIter = sortedNodes.begin();
	while (nodeIter != sortedNodes.end()) {
		FrameOperator* fromOperator = NULL;
		vector<Edge*> outEdges = graph.getOutputEdges(*nodeIter);
		vector<FrameAndLink>* fasVec = new vector<FrameAndLink>;

		m_LinkVectors.insert(fasVec);
		
		for  (size_t i = 0; i < outEdges.size(); i++) {
			auto_ptr<FrameLink> frameLink(NULL);
			FrameOperator* toOperator = NULL;
			
			if (outEdges[i] == NULL) {
				throw NullPtr();
			}

			frameLink.reset(new FrameLink(outEdges[i]->link));
			toOperator = nodeOpMap.find(outEdges[i]->to)->second;

			FrameAndLink fas;
			fas.link = frameLink.get();
			fas.frame = m_OpInputs.find(toOperator)->second;
			fasVec->push_back(fas);
			
			m_Links.insert(frameLink.release());
		}
		
		fromOperator = nodeOpMap.find(*nodeIter)->second;
		m_OpOutputLinks[fromOperator] = fasVec;
		nodeIter++;
	}
}

void GraphOperator::initReverse(const GraphNode& graphNode) {
	if (graphNode.getGraph() == NULL) {
		/* Todo: better error reporting */
		throw NullPtr();
	}
	const Graph& graph = *graphNode.getGraph();
	
	/* sort nodes in graph */
	list<const Node*> sortedNodes;
	topological_sort(graph, sortedNodes);
	sortedNodes.reverse();
	
	/* add operators and i/o frames from graph in order */
	map<const Node*, FrameOperator*> nodeOpMap;
	list<const Node*>::iterator nodeIter = sortedNodes.begin();
	
	while (nodeIter != sortedNodes.end()) {
		auto_ptr<FrameOperator> nodeFrameOp(NULL);
		auto_ptr<Frame> inputFrame(NULL);
		auto_ptr<Frame> outputFrame(NULL);
		
		nodeFrameOp = (*nodeIter)->createReverseOp();
		outputFrame.reset(new Frame((*nodeIter)->inputSpec()));
		inputFrame.reset(new Frame((*nodeIter)->outputSpec()));
		if (nodeFrameOp.get() == NULL) {
			/* TODO: better error reporting */
			throw NullPtr();
		}
		
		nodeOpMap[*nodeIter] = nodeFrameOp.get();
		
		m_OrderedFrameOps.push_back(nodeFrameOp.get());
		m_OpInputs[nodeFrameOp.get()] = inputFrame.get();
		m_OpOutputs[nodeFrameOp.get()] = outputFrame.get();
		
		m_FrameOps.insert(nodeFrameOp.release());
		m_Frames.insert(inputFrame.release());
		m_Frames.insert(outputFrame.release());
		nodeIter++;
	}
	
	/* take a second pass and add links. we couldn't add the links originally
	 * because they require all the nodes to be loaded. */
	nodeIter = sortedNodes.begin();
	while (nodeIter != sortedNodes.end()) {
		FrameOperator* fromOperator = NULL;
		vector<Edge*> outEdges = graph.getOutputEdges(*nodeIter);
		vector<FrameAndLink>* fasVec = new vector<FrameAndLink>;
		
		m_LinkVectors.insert(fasVec);
		
		for  (size_t i = 0; i < outEdges.size(); i++) {
			auto_ptr<FrameLink> frameLink(NULL);
			FrameOperator* toOperator = NULL;
			
			if (outEdges[i] == NULL) {
				throw NullPtr();
			}
			
			frameLink.reset(new FrameLink(outEdges[i]->link));
			toOperator = nodeOpMap.find(outEdges[i]->from)->second;
			
			FrameAndLink fas;
			fas.link = frameLink.get();
			fas.frame = m_OpInputs.find(toOperator)->second;
			fasVec->push_back(fas);
			
			m_Links.insert(frameLink.release());
		}
		
		fromOperator = nodeOpMap.find(*nodeIter)->second;
		m_OpOutputLinks[fromOperator] = fasVec;
		nodeIter++;
	}
}
